import torch
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from models.emotion_lstm import EmotionLSTM
from utils.dataset_loader import load_dataset
from config import *
import matplotlib.pyplot as plt


def train():
    # 加载数据集
    X_train, X_test, y_train, y_test = load_dataset(DATA_PATH, max_pad_len=MAX_PAD_LEN)

    # 转换为PyTorch张量
    X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
    y_train_tensor = torch.tensor(y_train, dtype=torch.long)

    # 创建数据加载器
    train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)

    model = EmotionLSTM(INPUT_DIM, HIDDEN_DIM, OUTPUT_DIM, NUM_LAYERS)  # 实例化模型类
    criterion = torch.nn.CrossEntropyLoss()  # 定义损失函数
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)  # 定义优化器

    # 用于绘图的记录变量
    train_losses = []
    train_accuracies = []

    # 训练模型
    for epoch in range(NUM_EPOCHS):
        model.train()  # 设置模型为训练模式
        total_loss = 0
        correct_predictions = 0
        total_samples = 0

        for batch_features, batch_labels in train_loader:
            optimizer.zero_grad()  # 梯度清零
            outputs = model(batch_features)  # 正向传播
            loss = criterion(outputs, batch_labels)  # 计算损失
            total_loss += loss.item()

            # 计算准确率
            _, predicted = torch.max(outputs, 1)  # 计算预测结果
            correct_predictions += (predicted == batch_labels).sum().item()  # 计算预测正确的样本数
            total_samples += batch_labels.size(0)  # 计算总样本数

            loss.backward()  # 反向传播
            optimizer.step()  # 更新模型参数

        epoch_loss = total_loss / len(train_loader)  # 计算当前epoch的平均损失
        epoch_accuracy = correct_predictions / total_samples  # 计算当前epoch的准确率

        train_losses.append(epoch_loss)
        train_accuracies.append(epoch_accuracy)

        print(f'Epoch [{epoch + 1}/{NUM_EPOCHS}], Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.4f}')

    # 保存模型
    torch.save(model.state_dict(), MODEL_PATH)  # 保存模型参数

    # 绘制损失和准确率曲线
    plt.figure(figsize=(12, 5))

    # 绘制损失曲线
    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='Training Loss')
    plt.title('Training Loss Curve')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')

    # 绘制准确率曲线
    plt.subplot(1, 2, 2)
    plt.plot(train_accuracies, label='Training Accuracy', color='orange')
    plt.title('Training Accuracy Curve')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')

    # 显示图表
    plt.tight_layout()
    # plt.show()
    plt.savefig('training_metrics.png')


if __name__ == '__main__':
    train()
